from torchvision.models import vgg16
from torchvision import transforms
import torch
import torch.nn as nn


class PerceptualLossFn(nn.Module):
    def __init__(self, weight_path, device):
        super().__init__()

        vgg = vgg16()
        vgg.load_state_dict(torch.load(weight_path))
        vgg.to(device)

        self.vgg_part1 = nn.Sequential(*list(vgg.features.children())[:3])
        self.vgg_part2 = nn.Sequential(*list(vgg.features.children())[3:8])
        self.vgg_part3 = nn.Sequential(*list(vgg.features.children())[8:15])
        self.vgg_part4 = nn.Sequential(*list(vgg.features.children())[15:27])

        self.lossfn = nn.MSELoss(reduction='mean')
        # imagenet数据集的均值和标准差
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])
    
    def forward(self, pred, target):
        pred = self.normalize(pred)
        target = self.normalize(target)

        p1 = self.vgg_part1(pred)
        p2 = self.vgg_part2(p1)
        p3 = self.vgg_part3(p2)
        p4 = self.vgg_part4(p3)

        with torch.no_grad():
            t1 = self.vgg_part1(target)
            t2 = self.vgg_part2(t1)
            t3 = self.vgg_part3(t2)
            t4 = self.vgg_part4(t3)
        
        loss = (self.lossfn(p1, t1) + self.lossfn(p2, t2) + self.lossfn(p3, t3) + self.lossfn(p4, t4)) / 4.
        
        return loss
    
if __name__ == '__main__':
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # lossfn = PerceptualLossFn('vgg16-397923af.pth', device)
    # print(lossfn)
    print(vgg16())